#!/usr/bin/env python
import re
import math
import shutil
import argparse
import sys
import os
import numpy as np
class vaspInpGen():
    def __init__(self,title):
        self.title=title
        self.cell=''
        self.coord=[]
        self.geom=''
        self.shrink=''
        self.kpoint=''
        self.basis=''
        self.scf=''
        self.dimension=3
        self.symGroup='1'
        self.opt=''
        self.dft=''
        self.atoms=0
        self.guess=''
        self.vacuumZ=10
        self.vacuumY=10
        self.potDir='/home/zhonc0a/vasp533/potPAW.52_PBE/'
        self.elements=dict(H=1,He=2,Li=3,Be=4,B=5,C=6,N=7,O=8,F=9,Na=11,Mg=12,Al=13,Si=14,P=15,S=16,Cl=17,Br=35)
        self.rev_elements=dict(zip(self.elements.values(),self.elements.keys()))
        
    def readcif(self,cifFileName):
        cifFile=open(cifFileName,'r')  
        # read coordinate and cell parameters from cif file         
        for line in cifFile:
            line=line.replace(')','').replace('(','')
            if re.search('_cell_length_a',line):
                self.cell_a=float(line.split()[1])
            if re.search('_cell_length_b',line):
                self.cell_b=float(line.split()[1])
            if re.search('_cell_length_c',line):
                self.cell_c=float(line.split()[1])
            if re.search('_cell_angle_alpha',line):
                self.cell_alpha=float(line.split()[1])       
            if re.search('_cell_angle_beta',line):
                self.cell_beta=float(line.split()[1])   
            if re.search('_cell_angle_gamma',line):
                self.cell_gamma=float(line.split()[1]) 
            if re.search('_symmetry_Int_Tables_number',line):
                self.symGroup=float(line.split()[1]) 
                if float(self.symGroup) != 1:
                    print 'the symGroup in cif file is not 1, exit now'
                    sys.exit()
            if re.search('-?\d+\.\d+\s+-?\d+\.\d+\s+-?\d+.\d+\s+',line):
                self.atoms=self.atoms+1
                #line=line.replace(')','').replace('(','')
                #print line
                coord=[self.elements[line.split()[1]],float(line.split()[2]),float(line.split()[3]),float(line.split()[4])]
                self.coord.append(coord)
        #convert cell paramter to cell vector
        self.cell_vect_a=[float(self.cell_a),0.0,0.0]
        self.cell_vect_b=[self.cell_b*math.cos(math.radians(self.cell_gamma)),self.cell_b*math.sin(math.radians(self.cell_gamma)),0.0]
        angle2Y=math.acos(math.cos(math.radians(float(self.cell_alpha)))*math.cos(math.pi/2-math.radians(float(self.cell_gamma))))
        cx=self.cell_c*math.cos(math.radians(float(self.cell_beta)))
        cy=self.cell_c*math.cos(angle2Y)
        cz=(self.cell_c**2-cx**2-cy**2)**0.5
        self.cell_vect_c=[cx,cy,cz]
		# check self.coord for negative number and number that big than one
        
        for coord in self.coord:
            for i in range(len(coord[1:])):
                j = i+1
                while coord[j] < 0:
                    coord[j] += 1
                while coord[j] > 1:
                    coord[j] -= 1
        #print np.matrix(self.coord)[:,1:]
        #print np.matrix([self.cell_vect_a,self.cell_vect_b,self.cell_vect_c])
        cart_coord = (np.matrix(self.coord)[:,1:] * np.matrix([self.cell_vect_a,self.cell_vect_b,self.cell_vect_c])).tolist()
        self.coord = [ self.coord[cart_coord.index(coord)][:1] + coord for coord in cart_coord]
        self.coordType='cartesian'
        #sys.exit()
        cifFile.close()
        
#read infomation from STRUC file generated by crystal
    def readSTRUC(self,strucFileName):
        strucFile=open(strucFileName,'r')
        readCell=0
        readCoord=0
        coords=[]
        #bohr convert to angstrom
        B2A=0.529177249
        self.dimension=3
        for line in strucFile:
            if '$cell vectors' in line:
                readCell=1
                continue
            if '$coordinates' in line:
                readCoord=1
                readCell=0
                continue
            if '$END' in line:
                readCoord=0
                readCell=0
                continue
            if readCell==1:
                self.cell_vect_a=[float(i)*B2A for i in line.split()]
                readCell+=1
                continue
            if readCell==2:
                self.cell_vect_b=[float(i)*B2A for i in line.split()]
                if max(map(float,self.cell_vect_b)) > 499:
                    self.dimension=self.dimension-1
                readCell+=1
                continue
            if readCell==3:
                self.cell_vect_c=[float(i)*B2A for i in line.split()]   
                if max(map(float,self.cell_vect_c)) > 499:
                    self.dimension=self.dimension-1
            if readCoord==1:
                coord=[int(line.split()[4]),float(line.split()[1])*B2A,float(line.split()[2])*B2A,float(line.split()[3])*B2A]
                coords.append(line)     
                self.coord.append(coord)
            
        self.coordType='cartesian'
        strucFile.close()
        
    def writePOSPOT(self):
        self.latticeConstant=max(self.cell_vect_a)
        sortedCoord={}
        coordX=[]
        coordY=[]
        coordZ=[]
        for i in self.coord:
            if i[0] not in sortedCoord:
                sortedCoord[i[0]]=[]
            sortedCoord[i[0]].append([i[1],i[2],i[3]])
            coordX.append(i[1])
            coordY.append(i[2])
            coordZ.append(i[3])
        if self.coordType=='cartesian': 
        #middleX=sum(coordX)/len(coordX)
            middleY=sum(coordY)/len(coordY)
            middleZ=sum(coordZ)/len(coordZ)
            #spanX=max(coordX)-min(coordX)
            spanY=max(coordY)-min(coordY)
            spanZ=max(coordZ)-min(coordZ)    
            if self.dimension==1:
                self.cell_vect_b=[0.0,spanY+self.vacuumY,0.0]
                self.cell_vect_c=[0.0,0.0,spanZ+self.vacuumZ]
            if self.dimension==2:    
                self.cell_vect_c=[0.0,0.0,spanZ+self.vacuumZ]
            deltaB=(spanY+self.vacuumY)/2-middleY
            deltaC=(spanZ+self.vacuumZ)/2-middleZ
        cellA=' '.join(["%.4f" % (i/self.latticeConstant) for i in self.cell_vect_a])
        cellB=' '.join(["%.4f" % (i/self.latticeConstant) for i in self.cell_vect_b])
        cellC=' '.join(["%.4f" % (i/self.latticeConstant) for i in self.cell_vect_c])

        #num_of_atomS is the number of atoms per atomic species
        num_of_atoms=''
        elemnt_type=''
        #coords is the scaled cartesian coordinates
        coords=''
        potcar=open('POTCAR','w')
        for i in sortedCoord:
            elemnt_type=elemnt_type+self.rev_elements[i]+' '
            num_of_atoms=num_of_atoms+str(len(sortedCoord[i]))+' '
            sourcePotcar=open(self.potDir+self.rev_elements[i]+'/POTCAR')
            shutil.copyfileobj(sourcePotcar,potcar)
            if self.coordType=='cartesian':
                for coord in sortedCoord[i]:
                    if self.dimension<=2:
                        coord[2]=coord[2]+deltaC
                        if self.dimension==1:
                            coord[1]=coord[1]+deltaB
                    coords=coords+'  {:>18.14f}  {:>18.14f}  {:>18.14f}\n'.format(*[i/self.latticeConstant for i in coord])
            if self.coordType=='Direct':
                 for coord in sortedCoord[i]:
                     coords=coords+'  '+'  '.join([str(i) for i in coord])+'\n'
        poscar=open('POSCAR','w')
        poscar.write(self.title+'\n')
        poscar.write(str(self.latticeConstant)+'\n')
        poscar.write(cellA+'\n')
        poscar.write(cellB+'\n')
        poscar.write(cellC+'\n')
        poscar.write(elemnt_type+'\n')
        poscar.write(num_of_atoms+'\n')
        poscar.write(self.coordType+'\n')
        poscar.write(coords)
        potcar.close()
        poscar.close()
    def writeKPOINTS(self):
        self.kpoint='1,1,1'
        kpoints=open('KPOINTS','w')       
        kpoints.write(self.title+'\n')
        kpoints.write('0\n')
        kpoints.write('Gamma\n')
        kpoints.write(' '.join(self.kpoint.split(',')))
        kpoints.close()
        
    #def writeINCAR(self):
parser=argparse.ArgumentParser(description='Generate POSCAR and POTCAR file for vasp')
parser.add_argument("-v","--vacuum",metavar="vacuum layer", help="Set the width of vacuum layer (angstrom), default:10A",  type=float,default=10)
parser.add_argument("-b","--basis",metavar="BASIS SET", help="Set the directory of basis set. Default: /home/zhonc0a/vasp533/potPAW.52_PBE/",  type=str,default='/home/zhonc0a/vasp533/potPAW.52_PBE/')  
#parser.add_argument("-k","--shrink",help="set the shrink(k points), e.g. '2,3,4,5,6 to benchmark different K points' Default: 2",  type=str,default='2') 
#parser.add_argument("-l","--levelshift",help="Set the level shift to help scf converge(1-10). Default: 0", type=str,default='0') 
#parser.add_argument("-m","--fmixing",help='Set the fmixing to help scf converge(1-50)%%. Default: 30', type=str,default='30')
parser.add_argument("-d","--dimension",help='Set the dimension for periodicity (1,2,3). Default: 3', type=int)
#parser.add_argument("-s","--symmetry",help='Set the label for symmetry group. Default: 1', type=str,default='1')
#parser.add_argument("-o","--opt",help='Set the paramters for geometry optimization (A:ATOMONLY,F:FULLOPTG,CE:CELLONLY,I:ITATOCEL,CV:CVOLOPT,False). Default: FULLOPTG',type=str,default='False')
#parser.add_argument("-e","--external",help='Use external geometry(.gui) file as input', type=str)
#parser.add_argument("-g","--guess",help='Read initial guess from previous run', action='store_true')
parser.add_argument("-D","--directory",help='Creat new directory with name filename-dirname, default dirname is empty, input None means not to generate directory.',default='',type=str)
#parser.add_argument("-v","--version",action='version', version='%(prog)s 2.0')
parser.add_argument("structureFile",nargs='+')
args=parser.parse_args()        
for structure in args.structureFile:
    title=''.join(re.split('[/.]',structure)[-2:-1])
    pospot=vaspInpGen(title)
    if structure.split('.')[-1].lower()=='cif':
        pospot.readcif(structure)
    if structure.split('.')[-1].lower()=='struc':
        pospot.readSTRUC(structure)
    pospot.vacuumZ=args.vacuum
    pospot.vacuumY=args.vacuum
    if args.dimension:
        pospot.dimension=args.dimension  
    pospot.potDir=args.basis
    if args.directory != 'None':
        current_dir=os.getcwd()
        dir_name=title + args.directory
        if not os.path.exists(dir_name):
            os.makedirs(dir_name)
        os.chdir(dir_name)
#print test.cell_vect_a
##print test.cell_vect_b
##print test.cell_vect_c
##print test.coord
        pospot.writePOSPOT()
        pospot.writeKPOINTS()
        os.chdir(current_dir)
    else:
        pospot.writePOSPOT()
        pospot.writeKPOINTS()
        

        
                
        
           
